from io import open
import glob
import os
import string
import unicodedata
import random
import time
import math
import torch
import torch.nn as nn
from torch.optim import Adam
from tqdm import tqdm


all_letters = string.ascii_letters + ".,;'"
data_path = '../../data/names/'
n_letters = len(all_letters)
# device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
device = torch.device('cpu')

#  去掉一些语言中的重音标记
def unicode_to_ascii(s):
    return ''.join(c for c in unicodedata.normalize('NFD', s)
                   if unicodedata.category(c) != 'Mn'
                   and c in all_letters)


def read_file(filename):
    lines = open(filename, encoding='utf-8').read().strip().split('\n')
    return [unicode_to_ascii(line) for line in lines]


def generate_category_line():
    category_lines = {}
    all_categories = []
    for filename in glob.glob(data_path + "*.txt"):
        category = os.path.splitext(os.path.basename(filename))[0]
        all_categories.append(category)
        lines = read_file(filename)
        category_lines[category] = lines
    return category_lines, all_categories


def line_to_tensor(line):
    tensor = torch.zeros(len(line), 1, n_letters, device=device)
    for li, letter in enumerate(line):
        tensor[li][0][all_letters.find(letter)] = 1
    return tensor



category_lines, all_categories = generate_category_line()
n_categories = len(all_categories)


class NameRNNModel(nn.Module):
    def __init__(self, input_size, hidden_size, output_size, num_layers=1):
        super(NameRNNModel, self).__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.rnn = nn.LSTM(input_size, hidden_size, num_layers)
        self.linear = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        rr, hn = self.rnn(x)
        return self.linear(rr)

    def init_hidden(self):
        return torch.zeros(self.num_layers, 1, self.hidden_size)


def category_from_output(output):
    top_n, top_i = output.topk(1)  # 返回values和indices，长度为指定的n。
    category_i = top_i[0].item()
    return all_categories[category_i], category_i


def generate_sample():
    category = random.choice(all_categories)
    line = random.choice(category_lines[category])
    category_tensor = torch.tensor([all_categories.index(category)], dtype=torch.long, device=device)
    line_tensor = line_to_tensor(line)
    return category, line, category_tensor, line_tensor


def time_since(since):
    now = time.time()
    s = now - since
    m = math.floor(s / 60)
    s -= m * 60
    return '%dm %ds' % (m, s)


criterion = nn.CrossEntropyLoss()
lr = 5e-3
input_size = n_letters
n_hidden = 128
output_size = n_categories
num_layers = 1
model = NameRNNModel(input_size, n_hidden, output_size, num_layers).to(device)
n_iter = 100000
print_every = 500
plot_every = 500
optimizer = Adam(model.parameters(), lr=lr)


def trainRNN(category_tensor, line_tensor):
    optimizer.zero_grad()
    output = model(line_tensor)
    output = output[-1::].squeeze(0)
    loss = criterion(output, category_tensor)
    loss.backward()
    optimizer.step()
    return output, loss.item()


def train_and_evaluate():
    all_losses = []
    start = time.time()
    current_loss = 0
    for it in tqdm(range(1, n_iter + 1)):
        category, line, category_tensor, line_tensor = generate_sample()
        output, loss = trainRNN(category_tensor, line_tensor)
        current_loss += loss
        if it % print_every == 0:
            predict, predict_i = category_from_output(output)
            correct = 'True' if predict == category else 'False'
            print('%d %d%% (%s) %.4f %s / %s %s' % (it, it/n_iter*100, time_since(start), loss, line, predict, correct))

        if it % plot_every == 0:
            all_losses.append(current_loss / plot_every)
            current_loss = 0

    with torch.no_grad():
        correct_num = 0
        for i in tqdm(range(10000)):
            category, line, category_tensor, line_tensor = generate_sample()
            output = model(line_tensor)[-1::].squeeze(0)
            predict, predict_i = category_from_output(output)
            if predict == category:
                correct_num += 1
        print(correct_num / 10000)

    return all_losses, int(time.time() - start)


if __name__ == '__main__':
    train_and_evaluate()



